import socket
import threading
import random
import base64
import os
from typing import Tuple


class UDPServer:
    """UDP file server supporting concurrent client downloads"""

    MIN_DATA_PORT = 50000
    MAX_DATA_PORT = 51000
    MAX_RETRIES = 5
    INITIAL_TIMEOUT = 1.0  # seconds
    MAX_BLOCK_SIZE = 1000  # bytes

    def __init__(self, port: int):
        """Initialize the server

        Args:
            port: Port number to listen on
        """
        self.port = port
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_socket.bind(('', port))
        print(f"[SERVER] Server started listening on port {port}")

    def _select_available_port(self) -> int:
        """Select an available port for data transfer

        Returns:
            Available port number

        Raises:
            OSError: If no available port is found
        """
        for _ in range(10):  # Try 10 times
            port = random.randint(self.MIN_DATA_PORT, self.MAX_DATA_PORT)
            try:
                test_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                test_socket.bind(('', port))
                test_socket.close()
                return port
            except OSError:
                continue
        raise OSError("Failed to find available port")

    def _send_response(self, sock: socket.socket, addr: Tuple[str, int], message: str):
        """Send response with error handling

        Args:
            sock: Socket to use
            addr: Destination address (host, port)
            message: Message to send
        """
        try:
            sock.sendto(message.encode('utf-8'), addr)
        except Exception as e:
            print(f"[ERROR] Failed to send response to {addr}: {e}")

    def handle_file_transmission(self, file_name: str, client_address: Tuple[str, int]):
        """Handle file transfer for a single client

        Args:
            file_name: Requested filename
            client_address: Client address (host, port)
        """
        client_socket = None
        try:
            # Create new socket for this client
            port = self._select_available_port()
            client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            client_socket.bind(('', port))
            client_socket.settimeout(self.INITIAL_TIMEOUT)

            print(f"[SERVER] Allocated port {port} for client {client_address}")

            file_path = os.path.join("server_files", file_name)
            if not os.path.exists(file_path):
                response = f"ERR {file_name} NOT_FOUND"
                self._send_response(self.server_socket, client_address, response)
                print(f"[SERVER] File {file_name} not found")
                return

            file_size = os.path.getsize(file_path)
            response = f"OK {file_name} SIZE {file_size} PORT {port}"
            print(f"[SERVER] Sending response: {response} to {client_address}")
            self._send_response(self.server_socket, client_address, response)

            with open(file_path, 'rb') as file:
                while True:
                    try:
                        data, addr = client_socket.recvfrom(2048)
                        request = data.decode('utf-8').strip()
                        print(f"[SERVER] Received request: {request} from {addr}")

                        if request.startswith(f"FILE {file_name} CLOSE"):
                            response = f"FILE {file_name} CLOSE_OK"
                            self._send_response(client_socket, addr, response)
                            print(f"[SERVER] File {file_name} transfer completed")
                            break

                        if request.startswith(f"FILE {file_name} GET"):
                            try:
                                parts = request.split()
                                start = int(parts[4])
                                end = int(parts[6])

                                # Validate request range
                                if start < 0 or end >= file_size or start > end:
                                    print(f"[ERROR] Invalid range request: {start}-{end}")
                                    continue

                                file.seek(start)
                                file_data = file.read(end - start + 1)
                                base64_data = base64.b64encode(file_data).decode('utf-8')
                                response = f"FILE {file_name} OK START {start} END {end} DATA {base64_data}"
                                self._send_response(client_socket, addr, response)
                            except (IndexError, ValueError) as e:
                                print(f"[ERROR] Error processing data request: {e}")
                                continue
                    except socket.timeout:
                        print(f"[WARNING] Client {addr} timeout")
                        continue
                    except Exception as e:
                        print(f"[ERROR] Error processing request: {e}")
                        break

        except Exception as e:
            print(f"[ERROR] File transfer error: {e}")
        finally:
            if client_socket:
                client_socket.close()

    def start(self):
        """Start the server main loop"""
        print("[SERVER] Server ready to accept requests...")
        while True:
            try:
                data, addr = self.server_socket.recvfrom(1024)
                request = data.decode('utf-8').strip()
                print(f"[SERVER] Received download request: {request} from {addr}")

                if not request.startswith("DOWNLOAD "):
                    print(f"[ERROR] Invalid request: {request}")
                    continue

                file_name = request.split()[1]
                threading.Thread(
                    target=self.handle_file_transmission,
                    args=(file_name, addr),
                    daemon=True
                ).start()

            except Exception as e:
                print(f"[ERROR] Server error: {e}")


if __name__ == "__main__":
    import sys

    if len(sys.argv) != 2:
        print("Usage: python UDPServer.py <port>")
        sys.exit(1)

    os.makedirs("server_files", exist_ok=True)
    server = UDPServer(int(sys.argv[1]))
    server.start()